from .base_exchange import *
from binance.spot import Spot  # 现货API
from binance.websocket.spot.websocket_stream import SpotWebsocketStreamClient  # 现货WebSocket - 新的类名
from binance.um_futures import UMFutures  # U本位合约
from binance.websocket.um_futures.websocket_client import UMFuturesWebsocketClient  # 合约WebSocket

import asyncio
import time

class BinanceSpotExchange(BaseExchange):
    """币安现货交易"""
    
    def __init__(self, api_key: str = "", api_secret: str = "", testnet: bool = False):
        super().__init__(api_key, api_secret)
        self.name = "BinanceSpot"
        self.testnet = testnet
        
        # 初始化REST API客户端
        base_url = 'https://testnet.binance.vision' if testnet else None
        self.client = Spot(
            api_key=api_key,
            api_secret=api_secret,
            base_url=base_url
        )
        
        # 初始化WebSocket客户端
        stream_url = "wss://testnet.binance.vision/stream" if testnet else "wss://stream.binance.com:9443/stream"
        self.ws = SpotWebsocketStreamClient(stream_url=stream_url)
        
        # 缓存数据
        self.tickers = {}
        self.orderbooks = {}
        self.balances = {}
        
    async def connect_ws(self):
        """连接WebSocket"""
        try:
            self.ws.start()
            
            # 订阅行情
            self.ws.mini_ticker(
                symbol="btcusdt",
                id=1,
                callback=self._handle_ticker
            )
            
            # 订阅深度
            self.ws.depth(
                symbol="btcusdt",
                id=2,
                level=20,
                speed=100,
                callback=self._handle_orderbook
            )
            
            # 订阅用户数据
            if self.api_key:
                listen_key = self.client.new_listen_key()['listenKey']
                self.ws.user_data(
                    listen_key=listen_key,
                    id=3,
                    callback=self._handle_user_data
                )
                asyncio.create_task(self._keep_alive_listen_key(listen_key))
                
        except Exception as e:
            raise Exception(f"Failed to connect websocket: {str(e)}")
            
    async def get_ticker(self, symbol: str) -> Ticker:
        """获取最新行情"""
        try:
            ticker = self.client.ticker_24hr(symbol)
            return Ticker(
                symbol=symbol,
                last_price=float(ticker['lastPrice']),
                bid_price=float(ticker['bidPrice']),
                ask_price=float(ticker['askPrice']),
                volume_24h=float(ticker['volume']),
                timestamp=ticker['closeTime']
            )
        except Exception as e:
            raise Exception(f"Failed to get ticker: {str(e)}")
            
    async def get_orderbook(self, symbol: str, limit: int = 20) -> OrderBook:
        """获取订单簿"""
        try:
            depth = self.client.depth(symbol, limit=limit)
            return OrderBook(
                asks=[(float(price), float(qty)) for price, qty in depth['asks']],
                bids=[(float(price), float(qty)) for price, qty in depth['bids']],
                timestamp=depth['lastUpdateId']
            )
        except Exception as e:
            raise Exception(f"Failed to get orderbook: {str(e)}")
            
    async def get_balance(self, asset: str = "") -> Dict[str, Balance]:
        """获取账户余额"""
        try:
            account = self.client.account()
            balances = {}
            for bal in account['balances']:
                if not asset or bal['asset'] == asset:
                    balances[bal['asset']] = Balance(
                        asset=bal['asset'],
                        free=float(bal['free']),
                        locked=float(bal['locked']),
                        total=float(bal['free']) + float(bal['locked']),
                        timestamp=time.time()
                    )
            return balances
        except Exception as e:
            raise Exception(f"Failed to get balance: {str(e)}")
            
    async def create_order(self, symbol: str, side: OrderSide, 
                         order_type: OrderType, quantity: float,
                         price: float = None) -> Order:
        """创建订单"""
        try:
            params = {
                "symbol": symbol,
                "side": side.value,
                "type": order_type.value,
                "quantity": quantity
            }
            
            if order_type == OrderType.LIMIT:
                params.update({
                    "timeInForce": "GTC",
                    "price": price
                })
                
            response = self.client.new_order(**params)
            
            return Order(
                exchange_id=self.name,
                order_id=str(response['orderId']),
                symbol=symbol,
                order_type=order_type,
                side=side,
                price=float(response['price']) if 'price' in response else 0,
                amount=float(response['origQty']),
                filled=float(response['executedQty']),
                status=self._convert_status(response['status']),
                timestamp=response['transactTime']
            )
            
        except Exception as e:
            raise Exception(f"Failed to create order: {str(e)}")
            
    def _handle_ticker(self, data: dict):
        """处理行情数据"""
        try:
            symbol = data['s']
            self.tickers[symbol] = Ticker(
                symbol=symbol,
                last_price=float(data['c']),
                bid_price=float(data['b']),
                ask_price=float(data['a']),
                volume_24h=float(data['v']),
                timestamp=data['E']
            )
        except Exception as e:
            self.logger.error(f"Error handling ticker: {str(e)}")
            
    def _handle_orderbook(self, data: dict):
        """处理深度数据"""
        try:
            symbol = data['s']
            self.orderbooks[symbol] = OrderBook(
                asks=[(float(p), float(q)) for p, q in data['a']],
                bids=[(float(p), float(q)) for p, q in data['b']],
                timestamp=data['E']
            )
        except Exception as e:
            self.logger.error(f"Error handling orderbook: {str(e)}")
            
    def _handle_user_data(self, data: dict):
        """处理用户数据"""
        try:
            event_type = data['e']
            
            if event_type == 'outboundAccountPosition':
                self._handle_account_update(data)
            elif event_type == 'executionReport':
                self._handle_order_update(data)
            elif event_type == 'balanceUpdate':
                self._handle_balance_update(data)
                
        except Exception as e:
            self.logger.error(f"Error handling user data: {str(e)}")
            
    async def _keep_alive_listen_key(self, listen_key: str):
        """保持listenKey有效"""
        while True:
            try:
                await asyncio.sleep(30 * 60)  # 每30分钟更新一次
                self.client.renew_listen_key(listen_key)
            except Exception as e:
                self.logger.error(f"Failed to renew listen key: {str(e)}")
                
    def _convert_status(self, status: str) -> OrderStatus:
        """转换订单状态"""
        status_map = {
            'NEW': OrderStatus.PENDING,
            'PARTIALLY_FILLED': OrderStatus.OPEN,
            'FILLED': OrderStatus.FILLED,
            'CANCELED': OrderStatus.CANCELED,
            'REJECTED': OrderStatus.REJECTED,
            'EXPIRED': OrderStatus.EXPIRED
        }
        return status_map.get(status, OrderStatus.PENDING)
        
    async def close(self):
        """关闭连接"""
        try:
            if self.ws:
                self.ws.stop()
                self.ws = None
        except Exception as e:
            self.logger.error(f"Error closing websocket: {str(e)}")
            
    def _validate_symbol(self, symbol: str):
        """验证交易对格式"""
        if not symbol or not isinstance(symbol, str):
            raise ValueError("Invalid symbol")
            
    def _validate_amount(self, amount: float):
        """验证数量"""
        if not amount or not isinstance(amount, (int, float)) or amount <= 0:
            raise ValueError("Invalid amount")
            
    def _validate_price(self, price: float):
        """验证价格"""
        if not price or not isinstance(price, (int, float)) or price <= 0:
            raise ValueError("Invalid price")
            
    async def create_futures_order(self, symbol: str, side: OrderSide, 
                                 order_type: OrderType, quantity: float,
                                 price: float = None, reduce_only: bool = False) -> Order:
        """创建期货订单"""
        try:
            raise Exception("This is not a futures account")
            
        except Exception as e:
            raise Exception(f"Failed to create futures order: {str(e)}")
            
    async def create_futures_market_order(self, symbol: str, side: OrderSide, 
                                        quantity: float, reduce_only: bool = False) -> Order:
        """创建期货市价单"""
        return await self.create_futures_order(
            symbol=symbol,
            side=side,
            order_type=OrderType.MARKET,
            quantity=quantity,
            reduce_only=reduce_only
        )
        
    async def create_futures_limit_order(self, symbol: str, side: OrderSide,
                                       price: float, quantity: float,
                                       reduce_only: bool = False) -> Order:
        """创建期货限价单"""
        return await self.create_futures_order(
            symbol=symbol,
            side=side,
            order_type=OrderType.LIMIT,
            quantity=quantity,
            price=price,
            reduce_only=reduce_only
        )

class BinanceFuturesExchange(BaseExchange):
    """币安U本位合约交易"""
    
    def __init__(self, api_key: str = "", api_secret: str = "", testnet: bool = False):
        super().__init__(api_key, api_secret)
        self.name = "BinanceFutures"
        self.testnet = testnet
        self.is_futures = True
        
        # 初始化REST API客户端
        base_url = 'https://testnet.binancefuture.com' if testnet else None
        self.client = UMFutures(
            key=api_key,
            secret=api_secret,
            base_url=base_url
        )
        
        # 初始化WebSocket客户端
        stream_url = "wss://fstream-testnet.binance.com/ws" if testnet else "wss://fstream.binance.com/ws"
        self.ws = UMFuturesWebsocketClient(stream_url=stream_url)
        
        # 缓存数据
        self.tickers = {}
        self.orderbooks = {}
        self.balances = {}
        self.positions = {}
        
    async def connect_ws(self):
        """连接WebSocket"""
        try:
            self.ws.start()
            
            # 订阅行情
            self.ws.mini_ticker(
                symbol="btcusdt",
                id=1,
                callback=self._handle_ticker
            )
            
            # 订阅深度
            self.ws.depth(
                symbol="btcusdt",
                id=2,
                level=20,
                speed=100,
                callback=self._handle_depth
            )
            
            # 订阅标记价格
            self.ws.mark_price(
                symbol="btcusdt",
                id=3,
                speed=1000,
                callback=self._handle_mark_price
            )
            
            # 订阅用户数据
            if self.api_key:
                listen_key = self.client.new_listen_key()['listenKey']
                self.ws.user_data(
                    listen_key=listen_key,
                    id=4,
                    callback=self._handle_user_data
                )
                asyncio.create_task(self._keep_alive_listen_key(listen_key))
                
        except Exception as e:
            raise Exception(f"Failed to connect websocket: {str(e)}")
            
    async def get_positions(self, symbol: str = "") -> Dict[str, Position]:
        """获取持仓信息"""
        try:
            positions = {}
            account = self.client.account()
            
            for position in account['positions']:
                if not symbol or position['symbol'] == symbol:
                    positions[position['symbol']] = Position(
                        symbol=position['symbol'],
                        size=float(position['positionAmt']),
                        entry_price=float(position['entryPrice']),
                        leverage=float(position['leverage']),
                        margin=float(position['initialMargin']),
                        unrealized_pnl=float(position['unrealizedProfit']),
                        liquidation_price=float(position['liquidationPrice']),
                        timestamp=time.time()
                    )
            return positions
        except Exception as e:
            raise Exception(f"Failed to get positions: {str(e)}")
            
    async def create_order(self, symbol: str, side: OrderSide, 
                         order_type: OrderType, quantity: float,
                         price: float = None, reduce_only: bool = False) -> Order:
        """创建订单"""
        try:
            params = {
                "symbol": symbol,
                "side": side.value,
                "type": order_type.value,
                "quantity": quantity,
                "reduceOnly": reduce_only
            }
            
            if order_type == OrderType.LIMIT:
                params.update({
                    "timeInForce": "GTC",
                    "price": price
                })
                
            response = self.client.new_order(**params)
            
            return Order(
                exchange_id=self.name,
                order_id=str(response['orderId']),
                symbol=symbol,
                order_type=order_type,
                side=side,
                price=float(response['price']) if 'price' in response else 0,
                amount=float(response['origQty']),
                filled=float(response['executedQty']),
                status=self._convert_status(response['status']),
                timestamp=response['updateTime']
            )
            
        except Exception as e:
            raise Exception(f"Failed to create order: {str(e)}")
            
    async def set_leverage(self, symbol: str, leverage: int) -> bool:
        """设置杠杆倍数"""
        try:
            response = self.client.change_leverage(
                symbol=symbol,
                leverage=leverage
            )
            return response['leverage'] == leverage
        except Exception as e:
            raise Exception(f"Failed to set leverage: {str(e)}")
            
    async def get_funding_rate(self, symbol: str) -> float:
        """获取资金费率"""
        try:
            response = self.client.mark_price(symbol=symbol)
            return float(response['lastFundingRate'])
        except Exception as e:
            raise Exception(f"Failed to get funding rate: {str(e)}")
            
    def _handle_ticker(self, data: dict):
        """处理行情数据"""
        try:
            symbol = data['s']
            self.tickers[symbol] = Ticker(
                symbol=symbol,
                last_price=float(data['c']),
                bid_price=float(data['b']),
                ask_price=float(data['a']),
                volume_24h=float(data['v']),
                timestamp=data['E']
            )
        except Exception as e:
            self.logger.error(f"Error handling ticker: {str(e)}")
            
    def _handle_depth(self, data: dict):
        """处理深度数据"""
        try:
            symbol = data['s']
            self.orderbooks[symbol] = OrderBook(
                asks=[(float(p), float(q)) for p, q in data['a']],
                bids=[(float(p), float(q)) for p, q in data['b']],
                timestamp=data['E']
            )
        except Exception as e:
            self.logger.error(f"Error handling depth: {str(e)}")
            
    def _handle_mark_price(self, data: dict):
        """处理标记价格数据"""
        try:
            symbol = data['s']
            if symbol in self.tickers:
                self.tickers[symbol].mark_price = float(data['p'])
                self.tickers[symbol].funding_rate = float(data['r'])
                self.tickers[symbol].next_funding_time = int(data['T'])
        except Exception as e:
            self.logger.error(f"Error handling mark price: {str(e)}")
            
    def _handle_user_data(self, data: dict):
        """处理用户数据"""
        try:
            event = data.get('e')
            
            if event == 'ACCOUNT_UPDATE':
                self._handle_account_update(data)
            elif event == 'ORDER_TRADE_UPDATE':
                self._handle_order_update(data)
            elif event == 'MARGIN_CALL':
                self._handle_margin_call(data)
                
        except Exception as e:
            self.logger.error(f"Error handling user data: {str(e)}")
            
    async def _keep_alive_listen_key(self, listen_key: str):
        """保持listenKey有效"""
        while True:
            try:
                await asyncio.sleep(30 * 60)  # 每30分钟更新一次
                self.client.renew_listen_key(listen_key)
            except Exception as e:
                self.logger.error(f"Failed to renew listen key: {str(e)}")
                
    def _convert_status(self, status: str) -> OrderStatus:
        """转换订单状态"""
        status_map = {
            'NEW': OrderStatus.PENDING,
            'PARTIALLY_FILLED': OrderStatus.OPEN,
            'FILLED': OrderStatus.FILLED,
            'CANCELED': OrderStatus.CANCELED,
            'REJECTED': OrderStatus.REJECTED,
            'EXPIRED': OrderStatus.EXPIRED
        }
        return status_map.get(status, OrderStatus.PENDING)
        
    async def close(self):
        """关闭连接"""
        try:
            if self.ws:
                self.ws.stop()
                self.ws = None
        except Exception as e:
            self.logger.error(f"Error closing websocket: {str(e)}")
            
    def _validate_symbol(self, symbol: str):
        """验证交易对格式"""
        if not symbol or not isinstance(symbol, str):
            raise ValueError("Invalid symbol")
            
    def _validate_amount(self, amount: float):
        """验证数量"""
        if not amount or not isinstance(amount, (int, float)) or amount <= 0:
            raise ValueError("Invalid amount")
            
    def _validate_price(self, price: float):
        """验证价格"""
        if not price or not isinstance(price, (int, float)) or price <= 0:
            raise ValueError("Invalid price")
            
    async def create_futures_order(self, symbol: str, side: OrderSide, 
                                 order_type: OrderType, quantity: float,
                                 price: float = None, reduce_only: bool = False) -> Order:
        """创建期货订单"""
        try:
            if not self.is_futures:
                raise Exception("This is not a futures account")
                
            params = {
                "symbol": symbol,
                "side": side.value,
                "type": order_type.value,
                "quantity": quantity,
                "reduceOnly": reduce_only
            }
            
            if order_type == OrderType.LIMIT:
                params.update({
                    "timeInForce": "GTC",
                    "price": price
                })
                
            # 使用futures的下单接口
            response = self.client.new_order(**params)
            
            return Order(
                exchange_id=self.name,
                order_id=str(response['orderId']),
                symbol=symbol,
                order_type=order_type,
                side=side,
                price=float(response['price']) if 'price' in response else 0,
                amount=float(response['origQty']),
                filled=float(response['executedQty']),
                status=self._convert_status(response['status']),
                timestamp=response['updateTime']
            )
            
        except Exception as e:
            raise Exception(f"Failed to create futures order: {str(e)}")
            
    async def create_futures_market_order(self, symbol: str, side: OrderSide, 
                                        quantity: float, reduce_only: bool = False) -> Order:
        """创建期货市价单"""
        return await self.create_futures_order(
            symbol=symbol,
            side=side,
            order_type=OrderType.MARKET,
            quantity=quantity,
            reduce_only=reduce_only
        )
        
    async def create_futures_limit_order(self, symbol: str, side: OrderSide,
                                       price: float, quantity: float,
                                       reduce_only: bool = False) -> Order:
        """创建期货限价单"""
        return await self.create_futures_order(
            symbol=symbol,
            side=side,
            order_type=OrderType.LIMIT,
            quantity=quantity,
            price=price,
            reduce_only=reduce_only
        ) 